{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33b1af1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = ['CUB', 'Derm7pt', 'RIVAL10']\n",
    "use_dataset = datasets[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02f84d33",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import torch\n",
    "from tqdm.auto import tqdm\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "notebook_dir = os.getcwd()\n",
    "project_root_path = os.path.dirname(notebook_dir)\n",
    "sys.path.insert(0, project_root_path)\n",
    "\n",
    "from src.config import PROJECT_ROOT, CUB_CONFIG, DERM7PT_CONFIG, RIVAL10_CONFIG  # noqa: E402\n",
    "import numpy as np  # noqa: E402"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58764a62",
   "metadata": {},
   "outputs": [],
   "source": [
    "if use_dataset == 'CUB':\n",
    "    config_dict = CUB_CONFIG\n",
    "    DATASET_PATH =  os.path.join(PROJECT_ROOT, 'output', 'CUB')\n",
    "elif use_dataset == 'Derm7pt':\n",
    "    config_dict = DERM7PT_CONFIG\n",
    "    DATASET_PATH =  os.path.join(PROJECT_ROOT, 'output', 'Derm7pt')\n",
    "else:\n",
    "    config_dict = RIVAL10_CONFIG\n",
    "    DATASET_PATH =  os.path.join(PROJECT_ROOT, 'output', 'RIVAL10')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b1b97ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_concepts = config_dict['N_TRIMMED_CONCEPTS']\n",
    "num_classes = config_dict['N_CLASSES']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7a3a2a2",
   "metadata": {},
   "source": [
    "# Load and Transform Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "951320be",
   "metadata": {},
   "outputs": [],
   "source": [
    "C_hat_train = np.load(os.path.join(DATASET_PATH, 'C_hat_sigmoid_train.npy'))\n",
    "one_hot_Y_train = np.load(os.path.join(DATASET_PATH, 'Y_train.npy'))\n",
    "\n",
    "C_hat_test = np.load(os.path.join(DATASET_PATH, 'C_hat_sigmoid_test.npy'))\n",
    "one_hot_Y_test = np.load(os.path.join(DATASET_PATH, 'Y_test.npy'))\n",
    "\n",
    "if use_dataset == 'Derm7pt':\n",
    "    C_hat_val = np.load(os.path.join(DATASET_PATH, 'C_hat_sigmoid_val.npy'))\n",
    "    one_hot_Y_val = np.load(os.path.join(DATASET_PATH, 'Y_val.npy'))\n",
    "\n",
    "    C_hat_train = np.concatenate((C_hat_train, C_hat_val), axis=0)\n",
    "    one_hot_Y_train = np.concatenate((one_hot_Y_train, one_hot_Y_val), axis=0)\n",
    "\n",
    "class_level_concepts = np.load(os.path.join(DATASET_PATH, 'class_level_concepts.npy'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "657421ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_train = np.argmax(one_hot_Y_train, axis=1)\n",
    "Y_test = np.argmax(one_hot_Y_test, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff09a63a",
   "metadata": {},
   "outputs": [],
   "source": [
    "C_train = []\n",
    "for y in Y_train:\n",
    "    C_train.append(class_level_concepts[y])\n",
    "\n",
    "C_train = np.array(C_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db9d6744",
   "metadata": {},
   "source": [
    "### Create DataLoaders for Prototypes Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5da141b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\")\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "# Create Dataloaders\n",
    "random_seed = 1\n",
    "\n",
    "def get_loaders(C_hat_train, C_hat_test):\n",
    "    if use_dataset == 'Derm7pt':\n",
    "        X_train = torch.tensor(C_hat_train, dtype=torch.float32)\n",
    "        Y_train = torch.tensor(one_hot_Y_train, dtype=torch.float32)\n",
    "    else:\n",
    "        X_train = torch.tensor(C_hat_train, dtype=torch.float32)\n",
    "        Y_train = torch.tensor(one_hot_Y_train, dtype=torch.float32)\n",
    "\n",
    "    X_test = torch.tensor(C_hat_test, dtype=torch.float32, device=device)\n",
    "    Y_test = torch.tensor(one_hot_Y_test, dtype=torch.float32, device=device)\n",
    "\n",
    "    # DATALOADERS\n",
    "    batch_size = 64\n",
    "    train_dataset = TensorDataset(X_train, Y_train)\n",
    "    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "    # test_dataset = TensorDataset(X_test, Y_test)\n",
    "    # test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    return train_loader, X_test, Y_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f46d4e63",
   "metadata": {},
   "source": [
    "# Add Noise to Concept Probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fbd4dd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_noise_to_concepts(c_hat, noise_level, random_seed=42):\n",
    "    np.random.seed(random_seed)\n",
    "    c_hat_noisy = c_hat.copy()\n",
    "\n",
    "    total_values = c_hat.size\n",
    "    values_to_modify = int(total_values * noise_level)\n",
    "\n",
    "    # Get flat indices for high and low values\n",
    "    flat_c_hat = c_hat.flatten()\n",
    "    all_indices = np.arange(len(flat_c_hat))\n",
    "    selected_indices = np.random.choice(all_indices, values_to_modify, replace=False)\n",
    "\n",
    "    # For clarity, we can still identify which selected values are high and low\n",
    "    selected_high = selected_indices[flat_c_hat[selected_indices] >= 0.5]\n",
    "    selected_low = selected_indices[flat_c_hat[selected_indices] < 0.5]\n",
    "\n",
    "    # high_indices = np.where(flat_c_hat >= 0.5)[0]\n",
    "    # low_indices = np.where(flat_c_hat < 0.5)[0]\n",
    "\n",
    "    # # Randomly select indices to modify\n",
    "    # high_count = min(values_to_modify // 2, len(high_indices))\n",
    "    # low_count = min(values_to_modify - high_count, len(low_indices))\n",
    "\n",
    "    # selected_high = np.random.choice(high_indices, high_count, replace=False)\n",
    "    # selected_low = np.random.choice(low_indices, low_count, replace=False)\n",
    "\n",
    "    # Apply noise using flat indexing\n",
    "    flat_noisy = c_hat_noisy.flatten()\n",
    "\n",
    "    # Push high values down toward 0.5\n",
    "    # high_values = flat_noisy[selected_high]\n",
    "    # max_decrease = high_values - 0.5\n",
    "    # noise = np.random.uniform(0, 1, len(selected_high)) * max_decrease\n",
    "    # flat_noisy[selected_high] -= noise\n",
    "    # flat_noisy[selected_high] = 0\n",
    "    flat_noisy[selected_high] = np.random.uniform(0, 0.5, len(selected_high))\n",
    "\n",
    "    # Push low values above 0.5\n",
    "    # low_values = flat_noisy[selected_low]\n",
    "    # max_increase = 0.5 - low_values\n",
    "    # noise = np.random.uniform(0, 1, len(selected_low)) * max_increase\n",
    "    # flat_noisy[selected_low] += noise\n",
    "    flat_noisy[selected_low] = np.random.uniform(0.5, 1, len(selected_low))\n",
    "\n",
    "    return flat_noisy.reshape(c_hat.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b930d66",
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_level = 0.1\n",
    "C_hat_test_noisy = add_noise_to_concepts(C_hat_test, noise_level)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d56f7ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "close_to_half = (np.sum((C_hat_test > 0.4) & (C_hat_test < 0.6)) / C_hat_test.size)\n",
    "close_to_half_noisy = (np.sum((C_hat_test_noisy > 0.4) & (C_hat_test_noisy < 0.6))/ C_hat_test.size)\n",
    "\n",
    "print(f\"{close_to_half*100}% of the original values are close to 0.5\")\n",
    "print(f\"{close_to_half_noisy*100}% of the original values are close to 0.5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2fa01de",
   "metadata": {},
   "source": [
    "# Test Phase 2 With Noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a946430",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "import os\n",
    "\n",
    "from src.models import PrototypeClassifier\n",
    "from src.training import train_epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a75b2224",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Global dictionary to store models in memory\n",
    "global_models = {\n",
    "    'logistic_regression': None,\n",
    "    'prototypes': None\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a486b8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_nearest_concept(instance, reference_concepts, reference_labels):\n",
    "    distances = np.sqrt(np.sum((reference_concepts - instance)**2, axis=1))\n",
    "    min_idx = np.argmin(distances)\n",
    "    return reference_labels[min_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d72a46e",
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperparams = {\n",
    "    \"CUB\": [50, 0.01, 0.01, 0.001],\n",
    "    \"Derm7pt\": [50, 0.01, 0.01, 0.001],\n",
    "    \"RIVAL10\": [50, 0.01, 0.01, 0.001],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d49c8b94",
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_levels = np.arange(0, 0.51, 0.05)\n",
    "n_repeats = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f5771f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize models in memory\n",
    "print(\"Training and storing models in memory for faster testing...\")\n",
    "\n",
    "# Create a dictionary to store all trained models\n",
    "trained_models = {}\n",
    "\n",
    "# Train and store all models\n",
    "lin_model = LogisticRegression(max_iter=1000)\n",
    "lin_model.fit(C_hat_train, Y_train)\n",
    "trained_models['lr'] = lin_model\n",
    "\n",
    "# Train prototype model\n",
    "train_loader, _, _ = get_loaders(C_hat_train, C_hat_test)\n",
    "proto_model = PrototypeClassifier(num_concepts, num_classes).to(device)\n",
    "num_epochs = hyperparams[use_dataset][0]\n",
    "optimizer = torch.optim.Adam(proto_model.parameters(), lr=hyperparams[use_dataset][1])\n",
    "lambda_binary = hyperparams[use_dataset][2]\n",
    "lambda_L1 = hyperparams[use_dataset][3]\n",
    "\n",
    "# Training loop for prototype model\n",
    "tqdm_loader = tqdm(range(num_epochs), desc=\"Training Prototypes\", leave=False)\n",
    "for epoch in tqdm_loader:\n",
    "    _, train_accuracy = train_epoch(proto_model, train_loader, optimizer, lambda_binary, lambda_L1, device=device)\n",
    "    tqdm_loader.set_postfix({\"Train Acc\": f\"{train_accuracy:.2f}%\"})\n",
    "\n",
    "trained_models['proto'] = proto_model\n",
    "print(\"All models trained and stored in memory\")\n",
    "\n",
    "# Create a fast evaluation function that uses the in-memory models\n",
    "def test_with_in_memory_models(C_hat_test, Y_test):\n",
    "    # Use the pre-trained models\n",
    "    results = {}\n",
    "\n",
    "    # Logistic Regression\n",
    "    results['lr'] = trained_models['lr'].score(C_hat_test, Y_test)\n",
    "\n",
    "    # Class-level concept nearest neighbor\n",
    "    # if C_train is not None:\n",
    "    #     correct_predictions = 0\n",
    "    #     for i, test_instance in enumerate(C_hat_test):\n",
    "    #         predicted_label = predict_nearest_concept(test_instance, C_train, Y_train)\n",
    "    #         true_label = Y_test[i]\n",
    "    #         if predicted_label == true_label:\n",
    "    #             correct_predictions += 1\n",
    "    #     results['class'] = correct_predictions / len(C_hat_test)\n",
    "\n",
    "    # Prototype model\n",
    "    proto_model = trained_models['proto']\n",
    "    proto_model.eval()\n",
    "    with torch.no_grad():\n",
    "        X_test = torch.tensor(C_hat_test, dtype=torch.float32, device=device)\n",
    "        Y_test_torch = torch.tensor(one_hot_Y_test, dtype=torch.float32, device=device)\n",
    "        real_labels = Y_test_torch.argmax(dim=1)\n",
    "        predictions = proto_model.predict(X_test)\n",
    "        results['proto'] = (predictions == real_labels).sum().item()/len(predictions)\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e45297d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "lr_results = []\n",
    "class_results = []\n",
    "proto_results = []\n",
    "\n",
    "# Run the noise level experiment with in-memory models\n",
    "start_time = time.time()\n",
    "\n",
    "for nl in tqdm(noise_levels, desc=\"Testing noise levels\"):\n",
    "    # print(f\"Testing noise of {nl*100:.0f}%\")\n",
    "    lr_acc, class_acc, proto_acc = [], [], []\n",
    "\n",
    "    # Create a progress bar for the repeats\n",
    "    for repeat in tqdm(range(n_repeats), desc=f\"Noise {nl*100:.0f}%\", leave=False):\n",
    "        C_hat_test_noisy = add_noise_to_concepts(C_hat_test, nl, random_seed=random_seed + repeat)\n",
    "        # Use the in-memory models instead of reloading from disk\n",
    "        results = test_with_in_memory_models(C_hat_test_noisy, Y_test)\n",
    "        lr_acc.append(results[\"lr\"])\n",
    "        # class_acc.append(results[\"class\"])\n",
    "        proto_acc.append(results[\"proto\"])\n",
    "    lr_results.append(np.mean(lr_acc))\n",
    "    # class_results.append(np.mean(class_acc))\n",
    "    proto_results.append(np.mean(proto_acc))\n",
    "\n",
    "end_time = time.time()\n",
    "print(f\"\\nTotal time for all noise levels: {end_time - start_time:.2f} seconds\")\n",
    "print(f\"Average time per noise level: {(end_time - start_time)/len(noise_levels):.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0360151b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Create a figure with appropriate size\n",
    "plt.figure(figsize=(6,6))\n",
    "\n",
    "# Plot each algorithm's results\n",
    "plt.plot(noise_levels, lr_results, marker='o', linestyle='--', label='LR', linewidth=3, markersize=10)\n",
    "# plt.plot(noise_levels, class_results, marker='D', linestyle='--', label='CLC-D')\n",
    "plt.plot(noise_levels, proto_results, marker='s', linestyle='-', label='CLPC', linewidth=3, markersize=10)\n",
    "\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.xlabel('Noise Level (%)', fontsize=16)\n",
    "plt.ylabel('Accuracy', fontsize=16)\n",
    "# plt.title(f'Model Performance vs. Noise Level ({use_dataset} Dataset)', fontsize=16)\n",
    "\n",
    "# Set y-axis limits from 0 to 1\n",
    "if use_dataset=='CUB':\n",
    "    plt.ylim(0.1, 0.8)\n",
    "elif use_dataset=='Derm7pt':\n",
    "    plt.ylim(0.25, 0.7)\n",
    "else:\n",
    "    plt.ylim(0.2, 1)\n",
    "\n",
    "\n",
    "# Format x-axis ticks as percentages\n",
    "plt.xticks(noise_levels, [f'{nl*100:.0f}' for nl in noise_levels], fontsize=16)\n",
    "plt.yticks(fontsize=16)\n",
    "\n",
    "# Add a horizontal line for random chance accuracy (1/num_classes)\n",
    "# random_chance = 1/num_classes\n",
    "# plt.axhline(y=random_chance, color='r', linestyle='--', alpha=0.7,\n",
    "#             label=f'Random Chance ({random_chance:.3f})')\n",
    "plt.legend(fontsize=16, loc='lower left')\n",
    "\n",
    "plt.tight_layout()\n",
    "output_dir = os.path.join(PROJECT_ROOT, 'output', 'plots')\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "plt.savefig(os.path.join(output_dir, f'{use_dataset}_noise.png'), dpi=600, bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ef5d859",
   "metadata": {},
   "outputs": [],
   "source": [
    "# proto_results"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "3.11.9",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
